open util/graph[Node] as graph
open util/ordering[State] as so
open util/ordering[W] as wo
open util/integer as integer

// Graphs with weighted edges:
sig Node { w: Node -> lone W }
// We start with one root node from which we create the complete spanning tree:
one sig Root extends Node {}
// Weights:
sig W { }
// The weight from edge x to y is the same as the weight from edge y to x:
fact wSymm { all x,y: Node | w[x,y] = w[y,x] }
// There is always one or more edges that have the highest possible weight:
fact wDiv { some x,y: Node | w[x,y]= wo/last }

// Define an edge in terms of weighted edges:
fun edge: Node -> set Node {
  { x,y: Node | some w[x,y] }
}
fact niceGraph {
  noSelfLoops[edge]
  undirected[edge] // edge relation is symmetric
  stronglyConnected[edge] // no lose parts
}

// States and initialisation as before:
sig State {
  // Set of tree nodes in current state:
  tree: set Node,
  // Parent pointers in current state:
  parent: Node -> lone Node
}

// Initializes the tree by setting a root node (a node without a parent).
pred init {
  let fs = so/first | { fs.tree = Root and no fs.parent }
}

// Extension now looks for the edge with minimum weight:
// (Note: lte means 'larger than or equals')
pred extend [pre, post: State] {
  some x,y: Node | x in pre.tree and 
                   y !in pre.tree and
                   y in x.edge and
                   all y':Node | {
                       y' in x.edge => lte[w[x,y],w[x,y']]
                     } and 
                   post.tree = pre.tree + y and
                   post.parent = pre.parent + y->x
}

//Rest as before:
fact createTree {
  init
  all s: State - so/last |
    let s' = so/next[s] | extend[s,s']
}
run {} for exactly 2 Node, 2 State, 1 W

assert STfound {
  so/last.tree = Node
  treeRootedAt[~(so/last.parent),Root]
}
check STfound for exactly 4 Node, 4 State, 3 W expect 0

pred testExtend [pre, post: State] {
  some x,y: Node | x in pre.tree and
                   y !in pre.tree and
                   post.tree = pre.tree + y and
                   post.parent = pre.parent + y->x
}

// generate all possible states
// check if the smallest state is indeed the smallest state \
// in the list of all possible states
assert isMinimalGraph {
  init
  all s: State - so/last |
    let s' = so/next[s] | testExtend[s,s']
}
check isMinimalGraph for exactly 4 Node, 4 State, 3 W expect 0

fun toInt [w: W]: Int {
  Int[#prevs[w] + 1]
}
fun addW [n: Int, w: W]: Int {
  add[n,toInt[w]]
}


